package protocol

import (
	"log"
	"net"
	"sync"
)

type Manager struct {
	listener net.Listener
	workers  map[string]net.Conn
	jobs     map[string]*Job
	jobQueue chan *Job
	mu       sync.Mutex
}

func NewManager(address string) (*Manager, error) {
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return nil, err
	}

	return &Manager{
		listener: listener,
		workers:  make(map[string]net.Conn),
		jobs:     make(map[string]*Job),
		jobQueue: make(chan *Job, 100), // 缓冲区大小可以根据需求调整
	}, nil
}

func (m *Manager) Start() error {
	log.Println("Manager started, waiting for connections")
	for {
		conn, err := m.listener.Accept()
		if err != nil {
			return err
		}
		log.Printf("New connection accepted from %s", conn.RemoteAddr())
		go m.handleConnection(conn)
	}
}

func (m *Manager) handleConnection(conn net.Conn) {
	log.Printf("Handling connection from %s", conn.RemoteAddr())
	for {
		msg := &Message{}
		if err := msg.Decode(conn); err != nil {
			log.Printf("Error decoding message: %v", err)
			continue
		}
		log.Printf("Received message from %s: %v", conn.RemoteAddr(), msg.Command)

		switch msg.Command {
		case READY:
			m.handleReady(conn)
		case GIVE:
			m.handleGive(conn)
		case DONE:
			m.handleDone(conn, msg)
		case STATUS:
			m.handleStatus(conn, msg)
		case CANCEL:
			m.handleCancel(conn, msg)
		default:
			log.Printf("Unknown command received: %d", msg.Command)
		}
	}
}

func (m *Manager) handleStatus(conn net.Conn, msg *Message) {
	jobID := msg.Headers["job_id"]
	m.mu.Lock()
	job, exists := m.jobs[jobID]
	m.mu.Unlock()

	statusMsg := &Message{
		Command: STATUS,
		Headers: map[string]string{"job_id": jobID},
	}

	if exists {
		statusMsg.Headers["status"] = job.Status
	} else {
		statusMsg.Headers["status"] = "unknown"
	}

	if err := statusMsg.Encode(conn); err != nil {
		log.Printf("Error sending status message: %v", err)
	}
}

func (m *Manager) handleCancel(conn net.Conn, msg *Message) {
	jobID := msg.Headers["job_id"]
	m.mu.Lock()
	delete(m.jobs, jobID)
	m.mu.Unlock()

	cancelMsg := &Message{
		Command: CANCEL,
		Headers: map[string]string{"job_id": jobID, "status": "cancelled"},
	}

	if err := cancelMsg.Encode(conn); err != nil {
		log.Printf("Error sending cancel message: %v", err)
	}
}

func (m *Manager) handleReady(conn net.Conn) {
	m.mu.Lock()
	defer m.mu.Unlock()

	addr := conn.RemoteAddr().String()
	m.workers[addr] = conn
}

func (m *Manager) handleGive(conn net.Conn) {
	log.Printf("Handling GIVE request from %s", conn.RemoteAddr())
	select {
	case job := <-m.jobQueue:
		log.Printf("Sending job %s to worker %s", job.ID, conn.RemoteAddr())
		takeMsg := &Message{
			Command: TAKE,
			Headers: map[string]string{"job_id": job.ID},
			Payload: job.Payload,
		}
		if err := takeMsg.Encode(conn); err != nil {
			log.Printf("Error sending TAKE message: %v", err)
			m.jobQueue <- job // 将作业放回队列
		}
	default:
		log.Printf("No jobs available, sending NO message to %s", conn.RemoteAddr())
		noMsg := &Message{Command: NO}
		if err := noMsg.Encode(conn); err != nil {
			log.Printf("Error sending NO message: %v", err)
		}
	}
}

func (m *Manager) AddJob(job *Job) {
	m.mu.Lock()
	m.jobs[job.ID] = job
	m.mu.Unlock()
	log.Printf("Added job %s to the queue", job.ID)
	m.jobQueue <- job
}

func (m *Manager) handleDone(conn net.Conn, msg *Message) {
	// Process the job result
	// You can implement job result handling logic here
}

func (m *Manager) Close() {
	m.listener.Close()
	for _, conn := range m.workers {
		conn.Close()
	}
}
